import torch
import numpy as np
from torch import nn
from torch.autograd import *
from torch import FloatTensor
from matplotlib import pyplot as plot

TIME_CYCLE = 25
BATCH_SIZE = 1
START = 0
END = 10


class TrigonoMetric(nn.Module):

    def __init__(self):
        super(TrigonoMetric, self).__init__()
        self.recurrent_layers = nn.RNN(
            input_size=1,
            hidden_size=32,
            num_layers=3,
            batch_first=True
        )
        self.linear = nn.Linear(32, 1)

    def forward(self, x, first_state=None):
        rnn_out, rnn_state = self.recurrent_layers(x, first_state)
        all_out = []
        for time_step in range(0, rnn_out.size(1)):
            rnn_out_x = rnn_out[:, time_step, :]
            all_out.append(rnn_out_x)
        return self.linear(torch.stack(all_out, 1))


def main():
    rnn = TrigonoMetric()
    optimizer = torch.optim.Adam(rnn.parameters(), 1e-2)
    loss_func = nn.MSELoss()
    draw_x = []
    draw_y = []
    for i in range(1000):
        seed_data = np.linspace(START, END, TIME_CYCLE)
        x = Variable(FloatTensor(np.sin(seed_data)[np.newaxis, :, np.newaxis]))
        y = Variable(FloatTensor(np.cos(seed_data)[np.newaxis, :, np.newaxis]))
        predict_y = rnn(x)
        loss = loss_func(predict_y, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if i % 20 == 0:
            plot.figure()
            plot.plot(seed_data, predict_y.data.numpy().reshape(TIME_CYCLE))
            loss_data = loss.data
            print("loss: %s" % str(loss_data))
            plot.show()


if __name__ == "__main__":
    main()
